#include <sys/time.h> 
#include <iostream>

#include "randomc.h"
#include "stocc.h"

#ifndef MULTIFILE_PROJECT
// If compiled as a single file then include these cpp files, 
// If compiled as a project then compile and link in these cpp files.
   #include "mersenne.cpp"             // code for random number generator
   #include "stoc1.cpp"                // random library source code
   #include "userintf.cpp"             // define system specific user interface
#endif

using namespace std;

void Yperm(float *Y, const int N)
{
  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  for (int n = 1; n < N; n++){
    int idx = sto.IRandom(0, n); // uniform over [0,n]
    float val = Y[n];
    Y[n] = Y[idx];
    Y[idx] = val;
  }
}

void Tstat(const float *Y, const float *X, float *T, const float *meanX, const float *sumquadX, const int N, const int K, const float meanY, const float sumquadY)
{
  float stat;
  for (int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;

    stat = 0.0;

    // compute test statistic
    for (int n = 0; n < N; n++)   
      stat += X[iter+n] * Y[n];

    stat = (stat-N*meanY*meanX[k])/sqrt(sumquadX[k]*sumquadY);
    T[k] = stat*sqrt(N-2.0)/sqrt(1.0-stat*stat);
  }
}

void TstatTmax(const float *Y, const float *X, float *T, const float *meanX, const float *sumquadX, const int N, const int K, const float meanY, const float sumquadY, float *Tmax)
{
  float zeta = 0.0;
  float stat;
  for (int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;

    stat = 0.0;

    // compute test statistic
    for (int n = 0; n < N; n++)   
      stat += X[iter+n] * Y[n];

    stat = (stat-N*meanY*meanX[k])/sqrt(sumquadX[k]*sumquadY);
    T[k] = stat*sqrt(N-2.0)/sqrt(1.0-stat*stat);

    if (zeta < fabs(T[k]))
      zeta = fabs(T[k]);
  }

  Tmax[0] = zeta;
}

void Pvalue(const float *T, const float *Ttilde, float *p, float *P, const float *Tmax, const int K)
{
  for(int k = 0; k < K; k++){
    if(fabs(Ttilde[k]) >= fabs(T[k]))
      p[k]++;
    if(Tmax[0] >= fabs(T[k]))
      P[k]++;
  }
}

long int seed = (long int)time(0); // global seed

int main(int argc, const char *argv[])
{
  // argv[1]: N
  // argv[2]: K
  // argv[3]: B

  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  int N = atoi(argv[1]);
  int K = atoi(argv[2]);
  int B = atoi(argv[3]);

  // print
  std::cout<<"(N="<<N<<",K="<<K<<",B="<<B<<")"<<"\n";

  // allocate memory  
  float *X;
  X = (float *)malloc(N*K*sizeof(float));
  float *Y;
  Y = (float *)malloc(N*sizeof(float));
  float *p;
  p = (float *)malloc(K*sizeof(float));
  float *P;
  P = (float *)malloc(K*sizeof(float));
  float *T;
  T = (float *)malloc(K*sizeof(float));
  float *Ttilde;
  Ttilde = (float *)malloc(K*sizeof(float));
  float *Tmax;
  Tmax = (float *)malloc(sizeof(float));
  float *meanX;
  meanX = (float *)malloc(K*sizeof(float));
  float *sumquadX;
  sumquadX = (float *)malloc(K*sizeof(float));

  // generate X and Y
  int count = 0;
  for (int n = 0; n < N; n++){
    Y[n] = sto.Normal(0, 1.0);
    for (int k = 0; k < K; k++){
      X[count] = sto.Normal(0, 1.0);
      count++;
    }
  }

  // start time profiling
  timeval t1, t2;
  double elapsed;
  gettimeofday(&t1, NULL);

  // compute mean of Y
  float meanY = 0.0;
  for (int n = 0; n < N; n++)
    meanY += Y[n];
  meanY /= N;

  for (int k = 0; k < K; k++){
    p[k] = 0;
    P[k] = 0;
    T[k] = 0.0;
    Ttilde[k] = 0.0;
  }

  // compute meanX
  for (int k = 0; k < K; k++){
    meanX[k] = 0.0;
    int iter = k*N;
    for (int n = 0; n < N; n++)
      meanX[k] += X[iter+n];
    meanX[k] /= N;
  }

  // compute sumquadX
  for (int k = 0; k < K; k++){
    sumquadX[k] = 0.0;
    int iter = k*N;
    for (int n = 0; n < N; n++)
      sumquadX[k] += (X[iter+n]-meanX[k]) * (X[iter+n]-meanX[k]);
  }

  // compute sumquadY
  float sumquadY = 0;
  for (int n = 0; n < N; n++)
    sumquadY += (Y[n]-meanY) * (Y[n]-meanY);
                                                                                                                          
  // start execution
  // compute T
  Tstat(Y, X, T, meanX, sumquadX, N, K, meanY, sumquadY);

  for (int b = 0; b < B; b++){
    // permute Y
    Yperm(Y, N);

    // compute test statistic and find maximum
    TstatTmax(Y, X, Ttilde, meanX, sumquadX, N, K, meanY, sumquadY, Tmax);

    // compute p-values
    Pvalue(T, Ttilde, p, P, Tmax, K);
  }

  gettimeofday(&t2, NULL);
  elapsed = (t2.tv_sec - t1.tv_sec) * 1000.0;      // sec to ms
  elapsed += (t2.tv_usec - t1.tv_usec) / 1000.0;

  std::cout<<"elapsed time:"<<elapsed<<" milliseconds"<<"\n";

  // cleanup
  free(Y);
  free(X);
  free(p);
  free(P);
  free(T);
  free(Ttilde);
  free(Tmax);
  free(meanX);
  free(sumquadX);

  return(0);
}
